#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include "parallel_for.h"

#define M 500
#define N 500

typedef struct {
    double **w;
    double **w_new;
    int i_start;
    int i_end;
    int j_start;
    int j_end;
} heat_args;

void* heat_func(int i, void *args) {
    heat_args *h_args = (heat_args*)args;
    for(int j = h_args->j_start; j < h_args->j_end; j++) {
        if(i > 0 && i < M-1 && j > 0 && j < N-1) {
            h_args->w_new[i][j] = (h_args->w[i-1][j] + h_args->w[i+1][j] + 
                                  h_args->w[i][j-1] + h_args->w[i][j+1]) / 4.0;
        }
    }
    return NULL;
}

int main(int argc, char *argv[]) {
    if(argc != 2) {
        printf("Usage: %s <num_threads>\n", argv[0]);
        return 1;
    }
    
    int num_threads = atoi(argv[1]);
    int iterations = 100;
    double diff = 0.0;
    
    // 初始化温度场
    double **w = (double **)malloc(M * sizeof(double *));
    double **w_new = (double **)malloc(M * sizeof(double *));
    for(int i = 0; i < M; i++) {
        w[i] = (double *)malloc(N * sizeof(double));
        w_new[i] = (double *)malloc(N * sizeof(double));
        for(int j = 0; j < N; j++) {
            w[i][j] = 0.0;
            if(i == 0 || j == 0 || i == M-1 || j == N-1) {
                w[i][j] = 100.0; // 边界温度
            }
        }
    }
    
    heat_args args = {w, w_new, 1, M-1, 1, N-1};
    
    clock_t start = clock();
    for(int iter = 0; iter < iterations; iter++) {
        // 并行更新温度场
        parallel_for(1, M-1, 1, heat_func, &args, num_threads);
        
        // 交换指针
        double **temp = w;
        w = w_new;
        w_new = temp;
    }
    clock_t end = clock();
    
    printf("Simulation completed in %f seconds\n", ((double)(end - start)) / CLOCKS_PER_SEC);
    
    // 释放内存...
    return 0;
}